// "Copyright [2021] <Copyright Shengkai Lin>"

#include "cluster/task_scheduler/task_generator/task_generator.h"

namespace ai_sim {

Define_Module(TaskGenerator);

void TaskGenerator::initialize() {
    EV_INFO << "INITing TaskGenerator" << endl;
    tmp_tasktype = TaskType(static_cast<int>(par("tmp_tasktype")));
    // error("tasktype %d", tmp_tasktype);
    maxnumGpus = par("maxnumGpus");
    maxnumBatches = par("maxnumBatches");
    modelsize = par("modelsize");
    compute_time = par("compute_time");
    task_interval = par("task_interval");
    numGpus = par("numGpus");
    numBatches = par("numBatches");
    numJobs = par("numJobs");
    total_tasks = 0;
    // omnetpp::cModule * controller_module =
    //     getParentModule()->getSubmodule("centralcontroller");
    // controller = dynamic_cast<CentralController*> (controller_module);
    // TaskMsg *a_msg = new TaskMsg;
    // a_msg->setKind(TASK_ALLOC_MSG);
    // a_msg->setNow_task(gen_task());
    // send(a_msg, "controller_port");
    // controller->alloc_task(gen_fully_utilized_task());
    int inputType = (
        getSimulation()->getSystemModule()->
        par("inputType")
    );
    if (inputType == 0) {
        scheduleAt(omnetpp::simTime(), new SelfMsg);
    }
    EV_INFO << "INIT TaskGenerator FINISHED" << endl;
}

Task TaskGenerator::gen_fully_utilized_task() {
    Task tmp_task(total_tasks,
                    tmp_tasktype,
                    maxnumGpus,
                    maxnumBatches,
                    modelsize,
                    compute_time);
    total_tasks++;
    EV<< "A MAX TASK GENERATED" << omnetpp::endl;
    return tmp_task;
}

Task TaskGenerator::gen_required_tasks() {
    Task tmp_task(total_tasks,
                    tmp_tasktype,
                    numGpus,
                    numBatches,
                    modelsize,
                    compute_time);
    total_tasks++;
    EV<< "A REQUIRED TASK GENERATED" << omnetpp::endl;
    return tmp_task;
}

void TaskGenerator::handleMessage(omnetpp::cMessage *msg) {
    if (msg->isSelfMessage() && total_tasks <numJobs) {
        TaskMsg *a_msg = new TaskMsg;
        a_msg->setKind(TASK_ALLOC_MSG);
        a_msg->setNow_task(gen_required_tasks());
        send(a_msg, "controller_port");
        // controller->alloc_task(gen_fully_utilized_task());
        scheduleAt(omnetpp::simTime()+task_interval, msg);
    } else {
        delete msg;
    }
}
}  // namespace ai_sim
